#!/usr/bin/python
# Disponivel no pacote de programas como: lhsfull.py
"""
Implements the Latin Hypercube Sampling technique as described by 
Iman and Conover, 1982, including correlation control
both for no correlation or for a specified correlation matrix 
for the sampled parameters
"""
from pylab import plot, figure,hist,show, subplot
import scipy.stats as stats
import numpy
from numpy.linalg import cholesky,inv

valid_dists = ['Normal','Triangular','Uniform']
def rank_restr(nparms=4, Iter=100, noCorrRestr=True, matCorr=None):
    """
    Retorna indices para a amostragem das variaveis com 
    a estrutura de correlacao desejada.
    """
    if noCorrRestr:
        x = [stats.randint.rvs(1,Iter+1,size=Iter) for i in xrange(nparms)]
    else:
        if matCorr==None:
            C=numpy.core.numeric.identity(nparms)
        else:
            C=numpy.matrix(matCorr)
        s0=numpy.arange(1.,Iter+1)/((Iter*1.)+1)
        s=stats.norm().ppf(s0)
        s1=[]
        for i in xrange(nparms):
            numpy.random.shuffle(s)
            s1.append(s.copy())
        S=numpy.matrix(s1)
        P=cholesky(C)
        Q=cholesky(numpy.corrcoef(S))
        Final=S.transpose()*inv(Q).transpose()*P.transpose()
        x = [stats.stats.rankdata(Final.transpose()[i,]) for i in xrange(nparms)]
        return x
        
def sample_dist(cum, dist='Normal', p=[0,1]):
    if dist == 'Normal':
        if len(p) == 2:
            d = stats.norm(p[0],1./p[1]).ppf(cum)
    elif dist=='Triangular':
        if len(p) ==3 and p[0]<=p[1]<=p[2]:
            scale=p[2]-p[0]
            d = stats.triang((float(p[1])-p[0])/scale,loc=p[0],scale=scale).ppf(cum)
    elif dist == 'Uniform':
        d = stats.uniform(p[0],p[1]-p[0]).ppf(cum)
    else:
        print '%s is an unsupported distribution!'%dist
        d=[]
    return d

def lhs(Pars, dists, parms, Iter=100, noCorrRestr=True, matCorr=None):
    ParsList=[]
    if len(Pars)==len(dists):
        indexes=rank_restr(nparms=len(dists), Iter=Iter, noCorrRestr=noCorrRestr, matCorr=matCorr)
        for i in xrange(len(Pars)):
            v=sample_dist(numpy.arange(0.0001,1.,1./Iter), dist=dists[i], p=parms[i])
            index=map(int,indexes[i]-1)
            ParsList.append(v[index])
    return tuple(ParsList)
            
if __name__=='__main__':
    cm = numpy.array([[1,0.4,0.3],
                      [0.4,1,0.3],
                      [0.3,0.3,1]])
    c=lhs(['Par1', 'Par2', 'Par3'],['Normal','Triangular','Uniform'], [[0,1], [1,5,8], [1,2]],1000, False,cm)
    subplot(131)
    hist(c[0], bins=30)
    subplot(132)
    hist(c[1], bins=30)
    subplot(133)
    hist(c[2], bins=30)
    print stats.spearmanr(c[0],c[1])
    print stats.spearmanr(c[0],c[2])
    show()
    
